# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# A script to run multinode training with submitit.
# --------------------------------------------------------

import argparse
import os
import uuid
from pathlib import Path

import main_linprobe as classification
import submitit


def parse_args():
    classification_parser = classification.get_args_parser()
    parser = argparse.ArgumentParser("Submitit for MAE linear probe", parents=[classification_parser])
    parser.add_argument("--ngpus", default=8, type=int, help="Number of gpus to request on each node")
    parser.add_argument("--nodes", default=2, type=int, help="Number of nodes to request")
    parser.add_argument("--timeout", default=4320, type=int, help="Duration of the job")
    parser.add_argument("--job_dir", default="", type=str, help="Job dir. Leave empty for automatic.")

    parser.add_argument("--partition", default="learnfair", type=str, help="Partition where to submit")
    parser.add_argument("--use_volta32", action='store_true', help="Request 32G V100 GPUs")
    parser.add_argument('--comment', default="", type=str, help="Comment to pass to scheduler")
    return parser.parse_args()


def get_shared_folder() -> Path:
    """
    获取一个共享文件夹的路径。首先检查 /checkpoint 目录是否存在，如果存在，
    则返回该路径下以当前用户名命名的文件夹路径。如果该文件夹不存在，则创建它。
    如果 /checkpoint 目录不存在，抛出异常。
    """
    user = os.getenv("USER")  # 获取当前操作系统中的用户名
    if Path("/checkpoint/").is_dir():  # 检查 /checkpoint 目录是否存在
        p = Path(f"/checkpoint/{user}/experiments")  # 生成共享文件夹路径
        p.mkdir(exist_ok=True)  # 创建该文件夹，如果已经存在则不报错
        return p  # 返回文件夹路径
    raise RuntimeError("No shared folder available")  # 如果 /checkpoint 目录不存在，则抛出异常


def get_init_file():
    """
    创建一个初始化文件，初始化文件的父目录需要存在，文件本身不能存在。
    文件名包含一个 UUID 以确保唯一性。每次调用该函数都会返回一个新的文件路径。
    """
    os.makedirs(str(get_shared_folder()), exist_ok=True)  # 创建共享文件夹（如果不存在）
    init_file = get_shared_folder() / f"{uuid.uuid4().hex}_init"  # 生成一个包含 UUID 的唯一文件路径
    if init_file.exists():  # 如果文件已经存在
        os.remove(str(init_file))  # 删除已有的文件
    return init_file  # 返回新生成的初始化文件路径


class Trainer(object):
    """
    训练器类，负责训练过程的管理，包括训练过程的配置、检查点恢复等功能。
    """
    def __init__(self, args):
        """
        初始化 Trainer 实例，保存训练参数。

        参数:
            args: 包含训练配置和超参数的命令行参数或字典。
        """
        self.args = args  # 保存传入的参数

    def __call__(self):
        """
        训练器的调用函数，实际上会调用分类训练脚本，并进行训练。
        """
        import main_linprobe as classification  # 导入线性探测分类脚本

        self._setup_gpu_args()  # 配置与 GPU 相关的参数
        classification.main(self.args)  # 调用分类训练脚本的 main 函数，开始训练

    def checkpoint(self):
        """
        检查点功能。如果有已有的检查点文件，则恢复训练；如果没有，则继续训练。
        返回一个延迟任务，等待执行器重新提交训练任务。
        """
        import os
        import submitit

        # 获取初始化文件路径，并设置分布式训练的 dist_url 参数
        self.args.dist_url = get_init_file().as_uri()
        checkpoint_file = os.path.join(self.args.output_dir, "checkpoint.pth")  # 检查点文件路径
        if os.path.exists(checkpoint_file):  # 如果检查点文件存在
            self.args.resume = checkpoint_file  # 设置恢复路径

        print("Requeuing ", self.args)  # 打印重新排队的参数信息
        empty_trainer = type(self)(self.args)  # 创建一个新的 Trainer 实例
        return submitit.helpers.DelayedSubmission(empty_trainer)  # 返回一个延迟任务，等待重新排队执行

    def _setup_gpu_args(self):
        """
        配置与 GPU 和分布式训练相关的参数，获取当前作业的 GPU 信息，并设置训练参数。
        """
        import submitit
        from pathlib import Path

        job_env = submitit.JobEnvironment()  # 获取提交作业的环境信息
        # 设置输出目录为包含作业 ID 的路径，确保每个作业的输出文件夹独立
        self.args.output_dir = Path(str(self.args.output_dir).replace("%j", str(job_env.job_id)))
        self.args.log_dir = self.args.output_dir  # 日志目录与输出目录相同
        self.args.gpu = job_env.local_rank  # 当前进程使用的 GPU 编号
        self.args.rank = job_env.global_rank  # 当前进程的全局编号（在分布式训练中使用）
        self.args.world_size = job_env.num_tasks  # 集群中的任务总数
        print(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")  # 打印任务组信息


def main():
    """
    主函数，用于解析参数并提交训练作业。通过 `submitit` 提交任务并配置作业资源。
    """
    args = parse_args()  # 解析命令行参数
    if args.job_dir == "":  # 如果没有指定 job_dir，则使用默认的共享文件夹路径
        args.job_dir = get_shared_folder() / "%j"

    # 使用作业 ID 生成一个独立的文件夹，方便追踪不同的实验
    executor = submitit.AutoExecutor(folder=args.job_dir, slurm_max_num_timeout=30)  # 创建执行器实例

    num_gpus_per_node = args.ngpus  # 每个节点的 GPU 数量
    nodes = args.nodes  # 节点数量
    timeout_min = args.timeout  # 超时时间（分钟）

    partition = args.partition  # SLURM 分区
    kwargs = {}  # 额外的参数
    if args.use_volta32:
        kwargs['slurm_constraint'] = 'volta32gb'  # 如果使用特定硬件（如 Volta 32GB GPU），设置约束
    if args.comment:
        kwargs['slurm_comment'] = args.comment  # 提交作业时的评论

    # 更新执行器的资源参数
    executor.update_parameters(
        mem_gb=40 * num_gpus_per_node,  # 每个节点的内存大小（GB）
        gpus_per_node=num_gpus_per_node,  # 每个节点的 GPU 数量
        tasks_per_node=num_gpus_per_node,  # 每个节点的任务数，通常一个任务对应一个 GPU
        cpus_per_task=10,  # 每个任务使用的 CPU 核心数
        nodes=nodes,  # 节点数
        timeout_min=timeout_min,  # 超时时间（分钟）
        # 以下是集群相关的 SLURM 配置
        slurm_partition=partition,  # SLURM 分区
        slurm_signal_delay_s=120,  # SLURM 信号延迟时间（秒）
        **kwargs  # 传递额外的配置参数
    )

    executor.update_parameters(name="mae")  # 设置作业名称为 "mae"

    args.dist_url = get_init_file().as_uri()  # 设置分布式训练的初始化文件 URL
    args.output_dir = args.job_dir  # 设置输出目录为指定的作业目录

    trainer = Trainer(args)  # 创建训练器实例
    job = executor.submit(trainer)  # 提交训练作业

    # 打印提交的作业 ID
    print(job.job_id)


if __name__ == "__main__":
    main()  # 调用主函数，开始执行
